import numpy as np
from numpy.typing import NDArray
from rich import print


class Solution:

    def __init__(self, *, guess_list: NDArray[np.integer],
                 scores: NDArray[np.integer]) -> None:
        '''
        guess_matrix :
            如下表示有3道题, 做了5次猜测
            [0, 3, 0],

            [2, 0, 2],

            [1, 1, 0],

            [1, 2, 3],

            [3, 0, 1]
        
        scores :
            如下表示5次猜测的得分
            [1, 1, 1, 0, 0]
        '''
        sorted_index = np.argsort(scores)
        self.guess_list = guess_list[sorted_index]
        self.scores = scores[sorted_index]

    def study_from_zero_to_one(self):
        '''
        根据初始化时的猜测和得分, 计算出可能的结果
        仅仅从得分为0和1的猜测中探索结果
        '''
        result = self._study_from_zero()
        result = self._study_from_one(result)
        return np.fromiter(result, dtype=np.ndarray)

    def _study_from_zero(self):
        '''
        返回值中每行代表一个题目的可能选项, 不足4项用 -1 补齐
        如：
            [0, 2, -1, -1]

            [0, 2, -1, -1]

            [1, 3, -1, -1]

            [0, 2, -1, -1]
        表示 第一题答案可能为 0,2 ;

            第二题答案可能为 0,2 ;

            第三题答案可能为 1,3 ;

            第四题答案可能为 0,2

        原理: 得分为0的猜测, 其中的选项肯定是错误答案, 由此筛选出可能的正确选项
        '''
        zero_guess = self.guess_list[self.scores == 0]
        zero_guess = zero_guess.transpose()
        if np.shape(zero_guess)[1] == 1:  # 第0维是题目数，第1维是测试数
            # 如果只有一次测试，不去重
            wrong = zero_guess
        else:
            wrong = np.unique(zero_guess, axis=0)

        result: list[NDArray[np.integer]] = []
        total = np.arange(0, 4)
        for w in wrong:
            t = np.setdiff1d(total, w)
            result.append(t)
        return result

    def _study_from_one(self, possible_list: list[NDArray[np.integer]]):
        '''
        传入study_from_zero的结果, 计算符合1分的猜测
        '''
        one_guess = self.guess_list[self.scores == 1]
        one_guess = one_guess.transpose()

        result = self._verify_from_one(one_guess, possible_list)
        return result

    def _verify_from_one(self,
                         test_guess: NDArray[np.integer],
                         possible_list: list[NDArray[np.integer]],
                         depth=0):
        # 对于第depth道题目
        possible_row = possible_list[depth]
        # 对于第p个备选项
        for p in possible_row:
            certian_list = possible_list.copy()
            certian_list[depth] = np.array([p])

            # 假设抽取的选项是可信的 逐行匹配答案
            test_bool_matrix = np.zeros((1, len(test_guess[0])))
            for i, certian_row in enumerate(certian_list):
                t = np.isin(test_guess[i], certian_row)
                test_bool_matrix = np.vstack((test_bool_matrix, t))

            # 在到最后一题前, 如果每个题目得分都有可能大于1，则假设成立，进入下一道题的判断
            test_score = np.sum(test_bool_matrix, axis=0)
            if depth != len(test_guess) - 1 and np.all(test_score >= 1):
                yield from self._verify_from_one(test_guess, certian_list,
                                                 depth + 1)
            elif depth == len(test_guess) - 1 and np.all(test_score == 1):
                yield np.concatenate(certian_list)

    def verify_from_n(self, n: int, answer: NDArray[np.integer]):
        '''
        验证给定答案是否符合给定得分
        '''
        n_guess = self.guess_list[self.scores == n]
        for g in n_guess:
            if self.check_answer(g, answer) != n:
                return False
        else:
            return True

    def check_answer(self, guess: NDArray[np.integer],
                     answer: NDArray[np.integer]) -> bool:
        return np.sum(np.equal(guess, answer))


solve = Solution(guess_list=np.array([[2, 1, 0], [0, 1, 3], [0, 3, 0],
                                      [2, 2, 3], [1, 3, 1]]),
                 scores=np.array([0, 1, 1, 1, 2]))

for answer in solve.study_from_zero_to_one():
    if solve.verify_from_n(2, answer):  # 2是正确答案
        print(answer)
